#ifndef _FULL_MULTIGRID_H
#define _FULL_MULTIGRID_H

#include "Element.h"
#include "FDM.h"
#include <chrono>
/**
 * @brief 加权Jacobi，Gauss-Seidel，超松弛，对称超松弛法
 * 
 */
enum Interator_Methon { WJ, GS, SOR, SSOR};
using std::cout;
using std::endl;


const double weight = 2.0 /3 ; // Jaccobi迭代和SOR迭代的松弛因子

double L2_Error(VectorXd v1, VectorXd v2) // 返回两个向量之间的L2距离
{
    double error = 0.0;
    for(int i =0; i < v1.size(); i++)
    {
        error += ((v1[i] - v2[i]) * (v1[i] - v2[i]));
    }
    return sqrt(error);
};

/**
 * @brief 对方程Au=F进行一次迭代
 * 
 * @param _A 方程系数矩阵
 * @param _u 迭代初始值 
 * @param _Rhs 方程右端项
 * @param _type 迭代方法
 * @return VectorXd 迭代一次后所得的近似解
 */
VectorXd Interator(SpMat &_A, VectorXd &_u, VectorXd &_Rhs, Interator_Methon _type )
{
    VectorXd v = Eigen::VectorXd::Zero(_Rhs.size());
    if(_type == WJ)
    {
        for( int i =0; i < _u.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i);it;++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -=(it.value() * (_u[i_row]));
                }
                v[i] = temp / a_ii;
            }
        }
        v = (1 - weight) * _u + weight * v ;
    }
    else if (_type == GS)
    {
        v = _u;
        for( int i =0; i < v.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i);it;++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -=(it.value() * (v[i_row]));
                }
                v[i] = temp / a_ii;
            }
        }
    }
    else if(_type == SOR)
    {
        v = _u;
        for( int i =0; i < v.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i);it;++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -=(it.value() * (v[i_row]));
                }
                v[i] = temp / a_ii * weight + v[i] * (1 - weight);
            }
        }
    }
    else if(_type == SSOR)
    {
        v = _u;
        for( int i =0; i < v.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i);it;++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -=(it.value() * (v[i_row]));
                }
                v[i] = temp / a_ii * weight + v[i] * (1 - weight);
            }
        }
        for( int i = v.size()-1; i >= 0; --i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i);it;++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -=(it.value() * (v[i_row]));
                }
                v[i] = temp / a_ii * weight + v[i] * (1 - weight);
            }
        }
    }
    return v ;
};

void DealWithBoundary(SpMat &A, VectorXd &Rhs, set<size_t> &boundary_index)
{
    int n_col = Rhs.size();
    VectorXd u = VectorXd::Zero(n_col);
    for(int i =0; i < n_col; i++)
    {
        if(boundary_index.find(i) != boundary_index.end())      //第i个是边界，则Rhs[i]/A(i,i)就是真解
        {
            u[i] = Rhs[i] / A.coeffRef(i,i);
        }
    }
    for(int i =0; i <n_col; i++)
    {
        if(boundary_index.find(i) != boundary_index.end())
        {
            for(Eigen::SparseMatrix<double>::InnerIterator it(A,i);it;++it)
            {
                int row = it.row();     //注意，以下所有的取A（i，j）的值需要保证该值存在，不然会使迭代器失效
                if(row == i)
                    continue;
                Rhs[row] -= A.coeffRef(row, i) * u[i];
                A.coeffRef(i,row) = 0.0;
                A.coeffRef(row,i) = 0.0;
            }
        }
    }
    for(int i =0; i < n_col; i++)
    {
        if(boundary_index.find(i) != boundary_index.end())      
        {
            Rhs[i] = u[i] *A.coeffRef(i,i);
        }
    }
};

class Multigrid_fullgrid
{
protected:
    vector<SpMat> Vec_Mat;  //保存所有的矩阵，Vec_Mat[0]为最细网格，Vec_Mat[n-1]为最粗网格
    int N_Layer;            //网格的层数
    vector<VectorXd> F_Rhs; //保存方程右端项
    vector<VectorXd> u_Lhs; //保存每一层网格方程的近似解
    vector<set<size_t> > Mesh_Boundary;   //每一层网格的边界点
    double Tolerance = 1e-8;   //求解使用的迭代法的终止精度
    size_t iter_num;        //迭代终止时整个过程所迭代的次数
    Interator_Methon interator_type = GS ;  //迭代的方法
    bool has_solved = false;            //判断是否以及解决
public:
    void set_boundary(vector<set<size_t> >&_Mesh_Boundary) { Mesh_Boundary = _Mesh_Boundary; };
    void set_matrix(vector<SpMat> &_Vec_Mat) { Vec_Mat = _Vec_Mat; };
    void set_rhs_f(vector<VectorXd> &_F_Rhs) { F_Rhs = _F_Rhs; };
    void set_tolerance(double _tolerance) { Tolerance = _tolerance; };
    void set_n_layer(int _n_layer) { N_Layer = _n_layer; };
    Multigrid_fullgrid( int _n) : N_Layer(_n), iter_num(0) { u_Lhs.resize(N_Layer); };
    Multigrid_fullgrid() = delete ;
    void compute();         //开始计算
    void compute(int _now_layer);       //计算第_now_layer层的结果
    void set_interator_type(Interator_Methon _interator_type ) { interator_type = _interator_type; };
    VectorXd get_res_vec() { return u_Lhs[0]; };
    VectorXd get_res_rhs() {return F_Rhs[0];};
    SpMat get_res_mat() { return Vec_Mat[0];};
    void V_cycle();                     //用v-cycle解方程
    void V_cycle_once(int _now_layer);          //做一次v-cycle     
    virtual VectorXd interpolation(VectorXd &u_2h) = 0 ;    //将粗网格下的值差值到细网格中
    virtual VectorXd restriction(VectorXd &u_h) = 0 ;       //将细网格下的值投影到粗网格中
};

void Multigrid_fullgrid::V_cycle()
{
    VectorXd u_init = VectorXd::Zero(F_Rhs[0].size());
    iter_num = 0 ;
    u_Lhs[0] = u_init;
    VectorXd u_iter1cycle = u_init;
    auto t_start = std::chrono::high_resolution_clock::now();
    while(!has_solved){
        V_cycle_once(0);
        if(L2_Error(u_Lhs[0],u_iter1cycle) < Tolerance)
            has_solved = true;
        u_iter1cycle = u_Lhs[0];
    }
    auto t_end = std::chrono::high_resolution_clock::now();
    float t_total = std::chrono::duration<float, std::milli>(t_end - t_start).count();
    cout <<"calculate over!\n" << "The number of iterator is :" << iter_num << endl;
    cout << "It takes " << t_total <<".ms" << endl;
           
};

void Multigrid_fullgrid::V_cycle_once(int _now_layer)
{
    if(_now_layer == N_Layer -1)        //当前是最粗网格情况，采用共轭梯度法求解矩阵
    {
        SpMat A = Vec_Mat[_now_layer];
        VectorXd &f = F_Rhs[_now_layer];
        DealWithBoundary(A,f,Mesh_Boundary[_now_layer]);
        Eigen::ConjugateGradient<SpMat> Solver_sparse;
        Solver_sparse.setTolerance(1e-16);
        Solver_sparse.compute(A);
        u_Lhs[_now_layer] = Solver_sparse.solve(f);
        return ;
    }
    else{
        int vis1 = 5;
        VectorXd u = u_Lhs[_now_layer];
        SpMat A = Vec_Mat[_now_layer];
        VectorXd f = F_Rhs[_now_layer];
        DealWithBoundary(A,f,Mesh_Boundary[_now_layer]);
        //先在当前网格迭代vis1次
        VectorXd u_iter1 ;
        // cout << "Now_layer is :" << _now_layer <<"\tThe L2 error is:" << endl;

        for(int i =0; i <vis1; i++)
        {
            u_iter1 = Interator(A, u, f, interator_type) ;
            iter_num ++ ;
            // cout << L2_Error(u,u_iter1) << endl;
            u = u_iter1 ;    
        }
        VectorXd r_h = f - A *u ;
        VectorXd r_2h = restriction(r_h);
        F_Rhs[_now_layer + 1] = r_2h ;
        VectorXd zero_u_2h = VectorXd::Zero(r_2h.size());
        u_Lhs[_now_layer + 1] = zero_u_2h;
        V_cycle_once(_now_layer +1);
        VectorXd e_h = interpolation(u_Lhs[_now_layer + 1]);
        u = u + e_h ;
        // cout << "In the layer" << _now_layer <<"\nAfter multigrid, the L2 error is: " << L2_Error(u,u_iter1) <<"\n"<< endl;
        for(int i =0; i < vis1; i++)
        {
            u_iter1 = Interator(A, u, f, interator_type) ;
            iter_num ++ ;
            u = u_iter1 ;    
        }
        u_Lhs[_now_layer] = u_iter1;
    }
}

void Multigrid_fullgrid::compute(int _now_layer)
{   
    if(_now_layer == N_Layer -1)        //当前是最粗网格情况，采用共轭梯度法求解矩阵
    {
        SpMat &A = Vec_Mat[_now_layer];
        VectorXd &f = F_Rhs[_now_layer];
        DealWithBoundary(A,f,Mesh_Boundary[_now_layer]);
        Eigen::ConjugateGradient<SpMat> Solver_sparse;
        Solver_sparse.setTolerance(1e-16);
        Solver_sparse.compute(A);
        u_Lhs[_now_layer] = Solver_sparse.solve(f);
        return ;
    }
    else{
        int vis1 = 5;
        VectorXd u = VectorXd::Zero(F_Rhs[_now_layer].size());
        SpMat &A = Vec_Mat[_now_layer];
        VectorXd &f = F_Rhs[_now_layer];
        DealWithBoundary(A,f,Mesh_Boundary[_now_layer]);
        //先在当前网格迭代vis1次
        VectorXd u_iter1 ;
        cout << "Now_layer is :" << _now_layer <<"\tThe L2 error is:" << endl;

        for(int i =0; i <vis1; i++)
        {
            u_iter1 = Interator(A, u, f, interator_type) ;
            iter_num ++ ;
            if(L2_Error(u,u_iter1) < Tolerance)
            {
                u_Lhs[_now_layer] = u_iter1;
                return ;
            }
            cout << L2_Error(u,u_iter1) << endl;
            u = u_iter1 ;    
        }
        VectorXd r_h = f - A *u ;
        VectorXd r_2h = restriction(r_h);
        F_Rhs[_now_layer + 1] = r_2h ;
        compute(_now_layer +1);
        VectorXd e_h = interpolation(u_Lhs[_now_layer + 1]);
        u = u + e_h ;
        u_iter1 = Interator(A, u, f, interator_type) ;
        cout << "In the layer" << _now_layer <<"\nAfter multigrid, the L2 error is: " << L2_Error(u,u_iter1) <<"\n"<< endl;
        while(1)
        {
            u_iter1 = Interator(A, u, f, interator_type) ;
            iter_num ++ ;
            if(L2_Error(u,u_iter1) < Tolerance)
            {
                cout << L2_Error(u,u_iter1) << endl;
                u_Lhs[_now_layer] = u_iter1;
                return ;
            }
            u = u_iter1 ;    
        }
        
    }
};

void Multigrid_fullgrid::compute()
{
    auto t_start = std::chrono::high_resolution_clock::now();
    compute(0);
    auto t_end = std::chrono::high_resolution_clock::now();
    float t_total = std::chrono::duration<float, std::milli>(t_end - t_start).count();
    cout <<"calculate over!\n" << "The number of iterator is :" << iter_num << endl;
    cout << "It takes " << t_total <<".ms" << endl;
}


class DIM1_Multigrid_fullgrid : public Multigrid_fullgrid
{
public:
    DIM1_Multigrid_fullgrid(int _n) : Multigrid_fullgrid(_n) { };
    VectorXd interpolation(VectorXd &u_2h);
    VectorXd restriction(VectorXd &u_h);
};

VectorXd DIM1_Multigrid_fullgrid::interpolation(VectorXd &u_2h)
{
    int n_size = u_2h.size();
    VectorXd u_h = VectorXd::Zero(2*n_size -1);
    for(int j =0; j < n_size -1; j++)
    {
        // if(j %2 == 0)
        //     u_h[j] = u_2h[j/2];
        // else
        //     u_h[j] = (u_2h[j/2] + u_2h[j/2 + 1]) / 2;
        u_h[2*j] = u_2h[j];
        u_h[2*j+1] = (u_2h[j] + u_2h[j+1]) /2 ;
    }
    return u_h;
};

VectorXd DIM1_Multigrid_fullgrid::restriction(VectorXd &u_h)
{
    int u_h_size = u_h.size();
    int u_2h_size = u_h_size / 2 + 1 ;
    VectorXd u_2h = VectorXd::Zero(u_2h_size);
    u_2h[0] = u_h[0];
    u_2h[u_2h_size - 1] = u_h[u_h_size - 1];
    for(int i = 1; i < u_2h_size - 1; i++)
        u_2h[i] = (u_h[2*i-1] + 2* u_h[2 *i] + u_h[2*i+1]) / 4;
    return u_2h;
};


/**
 * @brief 在这里所使用的网格以及为有限元方法生成的Q1网格
 * 
 */
class DIM2_Multigrid_fullgrid : public Multigrid_fullgrid
{
private:
    int n_pow_x;  // 1<<n_pow_x 代表x方向上分割的单元
    int n_pow_y;  // 1<<n_pow_y y方向上分割的单元数
public:
    DIM2_Multigrid_fullgrid(int _n_layer, int _n_pow_x, int _n_pow_y) : Multigrid_fullgrid(_n_layer), n_pow_x(_n_pow_x), n_pow_y(_n_pow_y) { };
    VectorXd interpolation(VectorXd &u_2h);
    VectorXd restriction(VectorXd &u_h);
    size_t coordinate2index(size_t _i, size_t _j, int _now_layer)
    {
        size_t nx = ((1 << n_pow_x) >> _now_layer);
        size_t ny = ((1 << n_pow_y) >> _now_layer);
        return ( _j + _i *(nx + 1));
    }
};


VectorXd DIM2_Multigrid_fullgrid::interpolation(VectorXd &u_2h)
{
    ///先找到在第几层
    size_t nx = 1 << n_pow_x;
    size_t ny = 1 << n_pow_y;
    int _now_layer = 0;
    for( ; _now_layer < N_Layer; _now_layer ++)
    {
        if(((nx >> _now_layer) + 1) * ((ny >> _now_layer) + 1) == u_2h.size())
            break;
    }
    size_t nx_2h = nx >> _now_layer;
    size_t ny_2h = ny >> _now_layer;
    VectorXd u_h = VectorXd::Zero(F_Rhs[_now_layer -1].size());
    for(size_t i =0; i <= nx_2h; i++)
        for(size_t j =0; j <= ny_2h; j++)
        {
            u_h[coordinate2index(2*i, 2*j,_now_layer-1)] = u_2h[coordinate2index(i,j,_now_layer)];
        }
    for(size_t i =0; i < nx_2h; i++)
        for(size_t j =0; j < ny_2h; j++)
        {
            u_h[coordinate2index(2*i+1,2*j,_now_layer-1)] = (u_2h[coordinate2index(i,j,_now_layer)] + u_2h[coordinate2index(i+1,j,_now_layer)]) / 2;
            u_h[coordinate2index(2*i,2*j+1,_now_layer-1)] = (u_2h[coordinate2index(i,j,_now_layer)] + u_2h[coordinate2index(i,j+1,_now_layer)]) / 2;
            u_h[coordinate2index(2*i+1,2*j+1,_now_layer-1)] = (u_2h[coordinate2index(i,j,_now_layer)] + u_2h[coordinate2index(i+1,j,_now_layer)] + u_2h[coordinate2index(i,j+1,_now_layer)] + u_2h[coordinate2index(i+1,j+1,_now_layer)]) / 4;
        }
    return u_h;
};

VectorXd DIM2_Multigrid_fullgrid::restriction(VectorXd &u_h)
{
    ///先找到在第几层
    size_t nx = 1 << n_pow_x;
    size_t ny = 1 << n_pow_y;
    int _now_layer = 0;
    for( ; _now_layer < N_Layer; _now_layer ++)
    {
        if(((nx >> _now_layer) + 1) * ((ny >> _now_layer) + 1) == u_h.size())
            break;
    }
    _now_layer ++ ;     //代表粗网格在第几层
    size_t nx_2h = nx >> _now_layer ;
    size_t ny_2h = ny >> _now_layer ;
    VectorXd u_2h = VectorXd::Zero((nx_2h+1)*(ny_2h+1));
    for(size_t i =0; i <= nx_2h; i++)
        for(size_t j =0; j <= ny_2h; j++)
        {
            if(Mesh_Boundary[_now_layer].find(coordinate2index(i,j,_now_layer)) != Mesh_Boundary[_now_layer].end())     //当该点为边界点时
                u_2h[coordinate2index(i,j,_now_layer)] = u_h[coordinate2index(2*i,2*j,_now_layer -1)];
            else{
                u_2h[coordinate2index(i,j,_now_layer)] = 1.0 / 16 *(u_h[coordinate2index(2*i-1,2*j-1,_now_layer -1)] + u_h[coordinate2index(2*i-1,2*j+1,_now_layer -1)] + u_h[coordinate2index(2*i+1,2*j-1,_now_layer -1)] + u_h[coordinate2index(2*i+1,2*j+1, _now_layer -1)]
                        + 2 *(u_h[coordinate2index(2*i,2*j-1,_now_layer -1)] + u_h[coordinate2index(2*i,2*j+1,_now_layer -1)] + u_h[coordinate2index(2*i-1,2*j,_now_layer -1)] + u_h[coordinate2index(2*i+1,2*j,_now_layer -1)]) + 4 *u_h[coordinate2index(2*i,2*j,_now_layer -1)]);
            }
        }
    return u_2h;
};

/**
 * @brief 采用多重网格法求解possion方程
 * 
 * @param _x_l 左端点
 * @param _x_r 右端点
 * @param _n_pow 2^_n_pow代表分割的段数
 * @param _n_layer 网格重数
 * @param f possion方程右端
 * @param u 边界解析式
 */
void Multigarid_Solution_Dim1(double _x_l, double _x_r, int _n_pow, int _n_layer, function<double(double)> f, function<double(double)> u )
{
    int n_size = (1 << _n_pow);
    DIM1_Multigrid_fullgrid solution(_n_layer);
    vector<SpMat> vec_Mat;
    vector<VectorXd> vec_f;
    vector<set<size_t> > vec_bound;
    for(int i =0; i < _n_layer; i++)
    {
        FDM_DIM1 Fdm_sol(n_size >> i, _x_l, _x_r, f, u);
        Fdm_sol.build();
        vec_Mat.push_back(Fdm_sol.get_matrix());
        vec_f.push_back(Fdm_sol.get_Frhs());
        vec_bound.push_back(Fdm_sol.get_boundary()); 
        
    }
    solution.set_boundary(vec_bound);
    solution.set_matrix(vec_Mat);
    solution.set_rhs_f(vec_f);
//    solution.compute();
    solution.V_cycle();
    VectorXd v_res = solution.get_res_vec();
    double h = (_x_r - _x_l) / n_size ;
    VectorXd v_val = VectorXd::Zero(n_size + 1);
    for(int i = 0; i <= n_size; i++)
        v_val[i] = u(_x_l + i*h);
    cout << "The L2_error of solution and real_res is: " << L2_Error(v_val ,v_res) << endl;
    SpMat res_mat = solution.get_res_mat();
    VectorXd res_rhs = solution.get_res_rhs();
    cout << "L2_error of r_h is" << L2_Error(res_mat*v_res , res_rhs) << endl;
}

/**
 * @brief 采用多重网格求解二维possion方程
 * 
 * @param _lb 左下角坐标
 * @param _rt 右上角坐标
 * @param _nx_pow 2^(_nx_pow)代表x方向上的分割数
 * @param _ny_pow 2^(_ny_pow)代表y方向上的分割数
 * @param _n_layer 网格的重数
 * @param f 右端项方程
 * @param u 边界解析式
 * @param _acc 积分精读
 */
void Multigarid_Solution_Dim2(vector<double> _lb, vector<double> _rt, int _nx_pow,int _ny_pow, int _n_layer, function<double(double,double)> f, function<double(double,double)> u, int _acc )
{
    int nx_size = (1 << _nx_pow);
    int ny_size = (1 << _ny_pow);
    DIM2_Multigrid_fullgrid solution(_n_layer, _nx_pow, _ny_pow);
    vector<SpMat> vec_Mat;
    vector<VectorXd> vec_f;
    vector<set<size_t> > vec_bound;
    for(int i =0; i < _n_layer; i++)
    {
        Q1_Rectangel_fem_solution Fem_sol(_lb, _rt, nx_size >> i, ny_size >> i, _acc, f ,u);
        Fem_sol.build();
        vec_Mat.push_back(Fem_sol.get_Mat());
        vec_f.push_back(Fem_sol.get_Rhs());
        vec_bound.push_back(Fem_sol.get_boundary()); 
    }
    solution.set_boundary(vec_bound);
    solution.set_matrix(vec_Mat);
    solution.set_rhs_f(vec_f);
    // solution.compute();
    solution.V_cycle();
    // cout <<vec_Mat[0] << endl;
    // cout << vec_f[0] <<"\n"<< endl;
    // cout << vec_f[1] << endl;

    VectorXd v_res = solution.get_res_vec();
    double hx = (_rt[0] - _lb[0]) / nx_size ;
    double hy = (_rt[1] - _lb[0]) / ny_size ;
    VectorXd v_val = VectorXd::Zero((nx_size + 1) * (ny_size + 1));
    for(int i = 0; i <= ny_size; i++)
        for (int j =0; j <= nx_size; j++)
            v_val[i*(nx_size + 1) + j] = u(_lb[0] +j *hx,_lb[1] + i*hy);
    cout << "The L2_error of solution and real_res is: " << L2_Error( v_val , v_res) << endl;
    SpMat res_mat = solution.get_res_mat();
    // cout << res_mat << endl;
    VectorXd res_rhs = solution.get_res_rhs();
    // cout <<res_rhs << endl ;
    // cout << "L2_error of r_h is" <<res_mat*v_res - res_rhs << endl;
};

#endif